# scan_all_cases.py
import three_smoke_modified as tsm
import math
import numpy as np
import pandas as pd
from contextlib import redirect_stdout
import io
import multiprocessing as mp
import time

# 固定速度
v_value = 121.09

# 定义所有无人机和导弹的组合
drones = ['FY1', 'FY2', 'FY3', 'FY4', 'FY5']
missiles = ['M1', 'M2', 'M3']


def evaluate_case(params):
    drone_name, missile_name, alpha_deg = params
    print(f"Testing {drone_name} vs {missile_name} at alpha_deg = {alpha_deg}")

    with redirect_stdout(io.StringIO()):  # 抑制输出
        total, drops, bursts, indiv_times = tsm.main(drone_name, missile_name, alpha_deg, v_value, save_excel=False)

    print(f"{drone_name} vs {missile_name} at alpha_deg {alpha_deg}: Total shield time = {total}")

    return {
        'drone': drone_name,
        'missile': missile_name,
        'alpha_deg': alpha_deg,
        'total_shield': total,
        'drop_times': drops,
        'burst_times': bursts,
        'indiv_times': indiv_times
    }


def find_best_alpha_for_case(drone_name, missile_name):
    print(f"Finding best alpha for {drone_name} vs {missile_name}")

    # 初始扫描：0到360度，步长10度
    alpha_deg_values = np.arange(0, 361, 10)
    params_list = [(drone_name, missile_name, alpha_deg) for alpha_deg in alpha_deg_values]

    # 使用多进程池
    with mp.Pool(processes=mp.cpu_count()) as pool:
        results = pool.map(evaluate_case, params_list)

    # 找到最大总掩蔽时间对应的角度
    total_shields = [r['total_shield'] for r in results]
    best_index = np.argmax(total_shields)
    best_alpha_deg = results[best_index]['alpha_deg']
    best_total = results[best_index]['total_shield']

    print(f"Best alpha_deg for {drone_name} vs {missile_name}: {best_alpha_deg}, Total shield time: {best_total}")

    # 细化扫描：最佳角度附近±10度，步长1度
    refined_alpha_deg_values = np.arange(best_alpha_deg - 10, best_alpha_deg + 11, 1)
    refined_params_list = [(drone_name, missile_name, alpha_deg) for alpha_deg in refined_alpha_deg_values]

    with mp.Pool(processes=mp.cpu_count()) as pool:
        refined_results = pool.map(evaluate_case, refined_params_list)

    # 找到细化后的最佳角度
    refined_total_shields = [r['total_shield'] for r in refined_results]
    best_refined_index = np.argmax(refined_total_shields)
    best_refined_alpha_deg = refined_results[best_refined_index]['alpha_deg']
    best_refined_total = refined_results[best_refined_index]['total_shield']

    print(
        f"Best refined alpha_deg for {drone_name} vs {missile_name}: {best_refined_alpha_deg}, Total shield time: {best_refined_total}")

    # 返回最佳结果
    return refined_results[best_refined_index]


if __name__ == '__main__':
    start_time = time.time()
    all_results = []

    # 对每种无人机和导弹的组合进行处理
    for drone in drones:
        for missile in missiles:
            result = find_best_alpha_for_case(drone, missile)
            all_results.append(result)

    # 创建结果DataFrame
    df = pd.DataFrame(all_results)

    # 保存到Excel文件
    df.to_excel('angle_total.xlsx', index=False)

    end_time = time.time()
    print(f"所有计算完成，总耗时: {end_time - start_time:.2f}秒")
    print("结果已保存到 angle_total.xlsx")